import os
import zipfile
import tempfile
import shutil
import re
from pathlib import Path
from typing import Dict, Optional
import datetime

class HLSProjectGenerator:
    """Generate complete HLS project from user inputs"""
    
    def __init__(self):
        pass
    
    def extract_top_function(self, code: str) -> str:
        """Extract the top function name from the synthesis code"""
        # Look for function declarations like: void function_name(...) or int function_name(...)
        pattern = r'(?:void|int|float|double|char|bool)\s+(\w+)\s*\('
        matches = re.findall(pattern, code)
        if matches:
            return matches[0]
        raise RuntimeError("No function declaration found in the synthesis code")
    
    def generate_project(self, synthesis_code: str, testbench_code: str = None, 
                        header_code: str = None, device_part: str = None, 
                        clock_period: int = 10) -> str:
        """
        Generate complete HLS project and return zip file path
        
        Args:
            synthesis_code: HLS synthesis code
            testbench_code: Testbench code (optional)
            header_code: Header code (optional)
            device_part: FPGA device part (optional)
            clock_period: Clock period in nanoseconds (default: 10)
            
        Returns:
            Path to generated zip file
        """
        # Extract top function name
        top_function = self.extract_top_function(synthesis_code)
        project_name = f"{top_function}_autogen"
        
        # Set default device part if not provided
        if not device_part:
            device_part = "xc7z020-clg484-1"
        
        # Create project directory in current working directory
        current_dir = os.getcwd()
        project_dir = os.path.join(current_dir, project_name)
        
        # Remove existing project directory if it exists
        if os.path.exists(project_dir):
            shutil.rmtree(project_dir)
        
        # Create new project directory
        os.makedirs(project_dir, exist_ok=True)
        
        try:
            # Generate all files in the project directory
            self._generate_synthesis_file(project_dir, top_function, synthesis_code)
            self._generate_header_file(project_dir, top_function, header_code)
            self._generate_testbench_file(project_dir, top_function, testbench_code)
            self._generate_run_py(project_dir, top_function, project_name, device_part, clock_period)
            self._generate_hls_config(project_dir, top_function, device_part, clock_period)
            self._generate_description_json(project_dir, top_function, project_name, device_part)
            self._generate_readme(project_dir, top_function, project_name)
            self._generate_run_hls_tcl(project_dir, top_function, project_name, device_part, clock_period)
            self._generate_run_vitis_sh(project_dir)
            
            # Create zip file for download
            zip_path = os.path.join(tempfile.gettempdir(), f"{project_name}.zip")
            with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
                for root, dirs, files in os.walk(project_dir):
                    for file in files:
                        file_path = os.path.join(root, file)
                        arcname = os.path.relpath(file_path, current_dir)
                        zipf.write(file_path, arcname)
            
            return zip_path
            
        except Exception as e:
            # If something goes wrong, clean up the project directory
            if os.path.exists(project_dir):
                shutil.rmtree(project_dir, ignore_errors=True)
            raise e
    
    def _generate_synthesis_file(self, project_dir: str, top_function: str, code: str):
        """Generate synthesis .cpp file"""
        filename = os.path.join(project_dir, f"{top_function}.cpp")
        
        # Use user-provided code as-is
        with open(filename, 'w', encoding='utf-8') as f:
            f.write(code)
    
    def _generate_header_file(self, project_dir: str, top_function: str, header_code: str = None):
        """Generate header .h file"""
        filename = os.path.join(project_dir, f"{top_function}.h")
        
        if header_code:
            # Use user-provided header code as-is
            content = header_code
        else:
            # Generate default header without copyright
            content = f"""#ifndef {top_function.upper()}_H
#define {top_function.upper()}_H

#define MAX_SIZE 16

// Function declaration
void {top_function}(int* in1, int* in2, int* out_r, int size, int rep_count);

#endif // {top_function.upper()}_H
"""
        
        with open(filename, 'w', encoding='utf-8') as f:
            f.write(content)
    
    def _generate_testbench_file(self, project_dir: str, top_function: str, testbench_code: str = None):
        """Generate testbench .cpp file"""
        filename = os.path.join(project_dir, f"{top_function}_test.cpp")
        
        if testbench_code:
            # Use user-provided testbench code as-is
            content = testbench_code
        else:
            # Generate default testbench without copyright
            content = f"""#include <algorithm>
#include <cstdio>
#include <random>
#include <vector>
#include <iomanip>
#include <iostream>
#include "{top_function}.h"

using std::default_random_engine;
using std::generate;
using std::uniform_int_distribution;
using std::vector;

int gen_random() {{
    static default_random_engine e;
    static uniform_int_distribution<int> dist(0, 10);
    return dist(e);
}}

void print(int* data, int columns, int rows) {{
    vector<int> out(columns * rows);
    for (int r = 0; r < 10; r++) {{
        for (int c = 0; c < 10; c++) {{
            std::cout << std::setw(4) << data[r * columns + c] << " ";
        }}
        std::cout << "…\\n";
    }}
    for (int r = 0; r < 10; r++) {{
        std::cout << "   … ";
    }}
    std::cout << "⋱\\n\\n";
}}

void verify(vector<int>& gold, vector<int>& output) {{
    for (int i = 0; i < (int)output.size(); i++) {{
        if (output[i] != gold[i]) {{
            std::cout << "Mismatch " << i << ": gold: " << gold[i] << " device: " << output[i] << "\\n";
            print(output.data(), 16, 16);
            exit(EXIT_FAILURE);
        }}
    }}
}}

int main() {{
    static const int columns = 16;
    static const int rows = 16;
    vector<int> A(columns * rows);
    vector<int> B(columns * rows);
    vector<int> C(columns * rows, 0);
    
    generate(begin(A), end(A), gen_random);
    generate(begin(B), end(B), gen_random);
    
    std::cout << "A:\\n";
    print(A.data(), columns, rows);
    
    std::cout << "B:\\n";
    print(B.data(), columns, rows);
    
    {top_function}(A.data(), B.data(), C.data(), columns, 1);
    
    std::cout << "TEST PASSED\\n\\n";
    
    return EXIT_SUCCESS;
}}
"""
        
        with open(filename, 'w', encoding='utf-8') as f:
            f.write(content)
    
    def _generate_run_py(self, project_dir: str, top_function: str, project_name: str, device_part: str, clock_period: int):
        """Generate run.py file"""
        filename = os.path.join(project_dir, "run.py")
        
        # Calculate frequency in MHz comment
        frequency = 1000 / clock_period
        
        content = f"""
import vitis
import os

cwd = os.getcwd()+'/'

# Initialize session
client = vitis.create_client()
client.set_workspace(path='./workspace')

# Delete the component if it already exists
if os.path.exists('./workspace/{project_name}'):
\tclient.delete_component(name='{project_name}')

# Create component. Create new config file in the component folder of the workspace
comp = client.create_hls_component(name='{project_name}', cfg_file = ['hls_config.cfg'], template = 'empty_hls_component')

# Get handle of config file, then programmatically set desired options
cfg_file = client.get_config_file(path = './workspace/{project_name}/hls_config.cfg')
cfg_file.set_value (                 key = 'part',                  value = '{device_part}') 
cfg_file.set_value (section = 'hls', key = 'syn.file',              value = cwd+'{top_function}.cpp')
cfg_file.set_values(section = 'hls', key = 'tb.file',               values = [cwd+'{top_function}_test.cpp'])
cfg_file.set_value (section = 'hls', key = 'syn.top',               value = '{top_function}')
cfg_file.set_value (section = 'hls', key = 'clock',                 value = '{clock_period}') # {frequency:.0f}MHz
cfg_file.set_value (section = 'hls', key = 'flow_target',           value = 'vitis')
cfg_file.set_value (section = 'hls', key = 'package.output.syn',    value = '0')
cfg_file.set_value (section = 'hls', key = 'package.output.format', value = 'xo')
cfg_file.set_value (section = 'hls', key = 'csim.code_analyzer',    value = '0')

# Run flow steps
comp = client.get_component(name='{project_name}')
comp.run(operation='C_SIMULATION')
comp.run(operation='SYNTHESIS')
comp.run(operation='CO_SIMULATION')
"""
        
        with open(filename, 'w', encoding='utf-8') as f:
            f.write(content)
    
    def _generate_hls_config(self, project_dir: str, top_function: str, device_part: str, clock_period: int):
        """Generate hls_config.cfg file"""
        filename = os.path.join(project_dir, "hls_config.cfg")
        
        content = f"""part={device_part}

[hls]
clock={clock_period}
flow_target=vivado
syn.file={top_function}.cpp
syn.top={top_function}
tb.file={top_function}_test.cpp
package.output.format=ip_catalog
package.output.syn=false
"""
        
        with open(filename, 'w', encoding='utf-8') as f:
            f.write(content)
    
    def _generate_description_json(self, project_dir: str, top_function: str, project_name: str, device_part: str):
        """Generate description.json file"""
        filename = os.path.join(project_dir, "description.json")
        
        description = f"HLS project for {top_function} function targeting {device_part}"
        
        content = f"""{{
  "name": "{project_name}",
  "description": "{description}",
  "device_part": "{device_part}",
  "hls": {{"config": "hls_config.cfg"}}
}}
"""
        
        with open(filename, 'w', encoding='utf-8') as f:
            f.write(content)
    
    def _generate_readme(self, project_dir: str, top_function: str, project_name: str):
        """Generate README file"""
        filename = os.path.join(project_dir, "README")
        
        content = f"""HLS Project: {project_name}
Generated on: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}

This is an HLS project for the {top_function} function.

System Requirements
==================
Vitis 2025.1 version must be properly installed and the following system environment variable must be set:

VITIS_PATH = C:\Xilinx\2025.1\Vitis\bin

Make sure the Vitis tools are properly installed and accessible from the command line.

Files Included in this Package
==============================
README
{top_function}.cpp
{top_function}.h
{top_function}_test.cpp
run_hls.tcl
run_vitis_commandline.sh
run.py
hls_config.cfg
description.json

Running the Design using Vitis HLS (edit run_hls.tcl to set $hls_exec and enable specific run steps)
=========================================================
vitis-run --mode hls --tcl run_hls.tcl

Running the design using Vitis v++ and vitis-run commands
=====================================================
./run_vitis_commandline.sh

Running the design using python script
=====================================================
vitis -s run.py
"""
        
        with open(filename, 'w', encoding='utf-8') as f:
            f.write(content)
    
    def _generate_run_hls_tcl(self, project_dir: str, top_function: str, project_name: str, device_part: str, clock_period: int):
        """Generate run_hls.tcl file"""
        filename = os.path.join(project_dir, "run_hls.tcl")
        
        content = f"""#
# Copyright 1986-2022 Xilinx, Inc. All Rights Reserved.
# Copyright 2022-2025 Advanced Micro Devices, Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Create a project
open_component -reset component_{project_name} -flow_target vivado

# Add design files
add_files {top_function}.cpp
# Add test bench & files
add_files -tb {top_function}_test.cpp 

# Set the top-level function
set_top {top_function}

# ########################################################
# Create a solution
# Define technology and clock rate
set_part  {{{device_part}}}
create_clock -period {clock_period}

# Set variable to select which steps to execute
set hls_exec 2

csim_design
# Set any optimization directives
# End of directives

if {{$hls_exec == 1}} {{
\t# Run Synthesis and Exit
\tcsynth_design
}} elseif {{$hls_exec == 2}} {{
\t# Run Synthesis, RTL Simulation and Exit
\tcsynth_design
\tcosim_design
}} elseif {{$hls_exec == 3}} {{ 
\t# Run Synthesis, RTL Simulation, RTL implementation and Exit
\tcsynth_design
\tcosim_design
\texport_design
}} else {{
\t# Default is to exit after running csynth
\tcsynth_design
}}

exit
"""
        
        with open(filename, 'w', encoding='utf-8') as f:
            f.write(content)
    
    def _generate_run_vitis_sh(self, project_dir: str):
        """Generate run_vitis_commandline.sh file"""
        filename = os.path.join(project_dir, "run_vitis_commandline.sh")
        
        content = """#!/bin/sh
: <<COMMENT
Script to run Vitis v++ and vitis-run commands for csim, csynth and cosim
COMMENT

HLS_CONFIG=hls_config.cfg
WORK_DIR=tmp

if [ $# -ne 1 ]; then
    echo "Usage: ./run_vitis_commandline.sh <csim(1)/csynth+cosim(2)>"
    exit
fi

selection=$1

case $selection in
    1)
        echo "Running C Simulation using vitis-run"
        vitis-run --mode hls --csim --config $HLS_CONFIG --work_dir $WORK_DIR
        ;;
    2)
        echo "Running C Synthesis using v++"
        v++ --compile --mode hls --config $HLS_CONFIG --work_dir $WORK_DIR

        echo "Running Co-Simulation using vitis-run"
        vitis-run --mode hls --cosim --config $HLS_CONFIG --work_dir $WORK_DIR
        ;;
    *)
        echo "Invalid input"
        exit
        ;;
esac
"""
        
        with open(filename, 'w', encoding='utf-8') as f:
            f.write(content) 